/**
 * 计算两个向量之间的欧几里得距离
 * @param {Array} vec1 - 第一个向量
 * @param {Array} vec2 - 第二个向量
 * @returns {number} 两个向量之间的欧几里得距离
 */
function euclideanDistance(vec1, vec2) {
  if (!Array.isArray(vec1) || !Array.isArray(vec2)) {
    throw new Error('Both arguments must be arrays')
  }
  if (vec1.length !== vec2.length) {
    throw new Error('Both vectors must have the same length')
  }
  let sum = 0
  for (let i = 0; i < vec1.length; i++) {
    sum += Math.pow(vec1[i] - vec2[i], 2)
  }
  return Math.sqrt(sum)
}

/**
 * 初始化聚类中心
 * @param {Array} data - 数据集
 * @param {number} k - 聚类的数量
 * @returns {Array} 初始化的聚类中心
 */
function initializeCentroids(data, k) {
  const centroids = []
  const indices = []
  while (centroids.length < k) {
    const index = Math.floor(Math.random() * data.length)
    if (!indices.includes(index)) {
      indices.push(index)
      centroids.push(data[index])
    }
  }
  return centroids
}

/**
 * 分配数据点到最近的聚类中心
 * @param {Array} data - 数据集
 * @param {Array} centroids - 聚类中心
 * @returns {Array} 每个数据点所属的聚类索引
 */
function assignClusters(data, centroids) {
  const clusters = []
  for (let i = 0; i < data.length; i++) {
    let minDistance = Infinity
    let closestCentroidIndex = 0
    for (let j = 0; j < centroids.length; j++) {
      const distance = euclideanDistance(data[i], centroids[j])
      if (distance < minDistance) {
        minDistance = distance
        closestCentroidIndex = j
      }
    }
    clusters.push(closestCentroidIndex)
  }
  return clusters
}

/**
 * 更新聚类中心
 * @param {Array} data - 数据集
 * @param {Array} clusters - 每个数据点所属的聚类索引
 * @param {number} k - 聚类的数量
 * @returns {Array} 更新后的聚类中心
 */
function updateCentroids(data, clusters, k) {
  const centroids = []
  for (let i = 0; i < k; i++) {
    const clusterPoints = []
    for (let j = 0; j < data.length; j++) {
      if (clusters[j] === i) {
        clusterPoints.push(data[j])
      }
    }
    if (clusterPoints.length > 0) {
      const centroid = []
      for (let dim = 0; dim < clusterPoints[0].length; dim++) {
        let sum = 0
        for (let p = 0; p < clusterPoints.length; p++) {
          sum += clusterPoints[p][dim]
        }
        centroid.push(sum / clusterPoints.length)
      }
      centroids.push(centroid)
    } else {
      // 如果某个聚类为空，重新初始化一个随机的聚类中心
      centroids.push(data[Math.floor(Math.random() * data.length)])
    }
  }
  return centroids
}

/**
 * K - Means 聚类算法
 * @param {Array} data - 数据集
 * @param {number} k - 聚类的数量
 * @param {number} maxIterations - 最大迭代次数
 * @returns {Array} 每个数据点所属的聚类索引
 */
function kmeans(data, k, maxIterations = 100) {
  let centroids = initializeCentroids(data, k)
  let clusters
  for (let i = 0; i < maxIterations; i++) {
    clusters = assignClusters(data, centroids)
    const newCentroids = updateCentroids(data, clusters, k)
    if (JSON.stringify(newCentroids) === JSON.stringify(centroids)) {
      break
    }
    centroids = newCentroids
  }
  return clusters
}

export default kmeans
